// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System.Diagnostics.CodeAnalysis;
using Microsoft.EntityFrameworkCore.Internal;

namespace Microsoft.EntityFrameworkCore.Query.Internal;

/// <summary>
///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
///     the same compatibility standards as public APIs. It may be changed or removed without notice in
///     any release. You should only use it directly in your code with extreme caution and knowing that
///     doing so can result in application failures when updating to a new Entity Framework Core release.
/// </summary>
public class QueryOptimizingExpressionVisitor : ExpressionVisitor
{
    private static readonly List<MethodInfo> SingleResultMethodInfos = new()
    {
        QueryableMethods.FirstWithPredicate,
        QueryableMethods.FirstWithoutPredicate,
        QueryableMethods.FirstOrDefaultWithPredicate,
        QueryableMethods.FirstOrDefaultWithoutPredicate,
        QueryableMethods.SingleWithPredicate,
        QueryableMethods.SingleWithoutPredicate,
        QueryableMethods.SingleOrDefaultWithPredicate,
        QueryableMethods.SingleOrDefaultWithoutPredicate,
        QueryableMethods.LastWithPredicate,
        QueryableMethods.LastWithoutPredicate,
        QueryableMethods.LastOrDefaultWithPredicate,
        QueryableMethods.LastOrDefaultWithoutPredicate
        //QueryableMethodProvider.ElementAtMethodInfo,
        //QueryableMethodProvider.ElementAtOrDefaultMethodInfo
    };

    private static readonly MethodInfo StringCompareWithComparisonMethod =
        typeof(string).GetRequiredRuntimeMethod(nameof(string.Compare), typeof(string), typeof(string), typeof(StringComparison));

    private static readonly MethodInfo StringCompareWithoutComparisonMethod =
        typeof(string).GetRequiredRuntimeMethod(nameof(string.Compare), typeof(string), typeof(string));

    private static readonly MethodInfo StartsWithMethodInfo =
        typeof(string).GetRequiredRuntimeMethod(nameof(string.StartsWith), typeof(string));

    private static readonly MethodInfo EndsWithMethodInfo =
        typeof(string).GetRequiredRuntimeMethod(nameof(string.EndsWith), typeof(string));

    private static readonly Expression ConstantNullString = Expression.Constant(null, typeof(string));

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitBinary(BinaryExpression binaryExpression)
    {
        var left = Visit(binaryExpression.Left);
        var right = Visit(binaryExpression.Right);

        if (binaryExpression.NodeType != ExpressionType.Coalesce
            && left.Type != right.Type
            && left.Type.UnwrapNullableType() == right.Type.UnwrapNullableType())
        {
            if (left.Type.IsNullableValueType())
            {
                right = Expression.Convert(right, left.Type);
            }
            else
            {
                left = Expression.Convert(left, right.Type);
            }
        }

        if (binaryExpression.NodeType == ExpressionType.Equal
            || binaryExpression.NodeType == ExpressionType.NotEqual)
        {
            var leftNullConstant = IsNullConstant(left);
            var rightNullConstant = IsNullConstant(right);
            if (leftNullConstant || rightNullConstant)
            {
                var nonNullExpression = leftNullConstant ? right : left;
                if (nonNullExpression is MethodCallExpression methodCallExpression
                    && methodCallExpression.Method.DeclaringType == typeof(Queryable)
                    && methodCallExpression.Method.IsGenericMethod
                    && methodCallExpression.Method.GetGenericMethodDefinition() is MethodInfo genericMethod
                    && SingleResultMethodInfos.Contains(genericMethod))
                {
                    var result = Expression.Call(
                        (methodCallExpression.Arguments.Count == 2
                            ? QueryableMethods.AnyWithPredicate
                            : QueryableMethods.AnyWithoutPredicate)
                        .MakeGenericMethod(methodCallExpression.Type),
                        methodCallExpression.Arguments);

                    return binaryExpression.NodeType == ExpressionType.Equal
                        ? Expression.Not(result)
                        : result;
                }
            }
        }

        return binaryExpression.Update(left, binaryExpression.Conversion, right);
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitConditional(ConditionalExpression conditionalExpression)
    {
        var test = Visit(conditionalExpression.Test);
        var ifTrue = Visit(conditionalExpression.IfTrue);
        var ifFalse = Visit(conditionalExpression.IfFalse);

        if (ifTrue.Type != ifFalse.Type
            && ifTrue.Type.UnwrapNullableType() == ifFalse.Type.UnwrapNullableType())
        {
            if (ifTrue.Type.IsNullableValueType())
            {
                ifFalse = Expression.Convert(ifFalse, ifTrue.Type);
            }
            else
            {
                ifTrue = Expression.Convert(ifTrue, ifFalse.Type);
            }

            return Expression.Condition(test, ifTrue, ifFalse);
        }

        return conditionalExpression.Update(test, ifTrue, ifFalse);
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override ElementInit VisitElementInit(ElementInit elementInit)
    {
        var arguments = new Expression[elementInit.Arguments.Count];
        for (var i = 0; i < arguments.Length; i++)
        {
            arguments[i] = MatchExpressionType(
                Visit(elementInit.Arguments[i]),
                elementInit.Arguments[i].Type);
        }

        return elementInit.Update(arguments);
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitLambda<T>(Expression<T> lambdaExpression)
    {
        var body = Visit(lambdaExpression.Body);

        return body.Type != lambdaExpression.Body.Type
            ? Expression.Lambda(Expression.Convert(body, lambdaExpression.Body.Type), lambdaExpression.Parameters)
            : lambdaExpression.Update(body, lambdaExpression.Parameters);
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitMember(MemberExpression memberExpression)
    {
        var expression = memberExpression.Expression != null
            ? MatchExpressionType(
                Visit(memberExpression.Expression),
                memberExpression.Expression.Type)
            : null;

        var visitedExpression = memberExpression.Update(expression);

        return TryOptimizeMemberAccessOverConditional(visitedExpression) ?? visitedExpression;
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override MemberAssignment VisitMemberAssignment(MemberAssignment memberAssignment)
    {
        var expression = MatchExpressionType(
            Visit(memberAssignment.Expression),
            memberAssignment.Expression.Type);

        return memberAssignment.Update(expression);
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitMethodCall(MethodCallExpression methodCallExpression)
    {
        if (Equals(StartsWithMethodInfo, methodCallExpression.Method)
            || Equals(EndsWithMethodInfo, methodCallExpression.Method))
        {
            if (methodCallExpression.Arguments[0] is ConstantExpression constantArgument
                && constantArgument.Value is string stringValue
                && stringValue == string.Empty)
            {
                // every string starts/ends with empty string.
                return Expression.Constant(true);
            }

            var newObject = Visit(methodCallExpression.Object)!;
            var newArgument = Visit(methodCallExpression.Arguments[0]);

            var result = Expression.AndAlso(
                Expression.NotEqual(newObject, ConstantNullString),
                Expression.AndAlso(
                    Expression.NotEqual(newArgument, ConstantNullString),
                    methodCallExpression.Update(newObject, new[] { newArgument })));

            return newArgument is ConstantExpression
                ? result
                : Expression.OrElse(
                    Expression.Equal(
                        newArgument,
                        Expression.Constant(string.Empty)),
                    result);
        }

        if (methodCallExpression.Method.IsGenericMethod
            && methodCallExpression.Method.GetGenericMethodDefinition() is MethodInfo methodInfo
            && (methodInfo.Equals(EnumerableMethods.AnyWithPredicate) || methodInfo.Equals(EnumerableMethods.All))
            && methodCallExpression.Arguments[0].NodeType is ExpressionType nodeType
            && (nodeType == ExpressionType.Parameter || nodeType == ExpressionType.Constant)
            && methodCallExpression.Arguments[1] is LambdaExpression lambda
            && TryExtractEqualityOperands(lambda.Body, out var left, out var right, out var negated)
            && (left is ParameterExpression || right is ParameterExpression))
        {
            var nonParameterExpression = left is ParameterExpression ? right : left;

            if (methodInfo.Equals(EnumerableMethods.AnyWithPredicate)
                && !negated)
            {
                var containsMethod = EnumerableMethods.Contains.MakeGenericMethod(methodCallExpression.Method.GetGenericArguments()[0]);
                return Expression.Call(null, containsMethod, methodCallExpression.Arguments[0], nonParameterExpression);
            }

            if (methodInfo.Equals(EnumerableMethods.All) && negated)
            {
                var containsMethod = EnumerableMethods.Contains.MakeGenericMethod(methodCallExpression.Method.GetGenericArguments()[0]);
                return Expression.Not(Expression.Call(null, containsMethod, methodCallExpression.Arguments[0], nonParameterExpression));
            }
        }

        if (methodCallExpression.Method.IsGenericMethod
            && methodCallExpression.Method.GetGenericMethodDefinition() is MethodInfo containsMethodInfo
            && containsMethodInfo.Equals(QueryableMethods.Contains))
        {
            var typeArgument = methodCallExpression.Method.GetGenericArguments()[0];
            var anyMethod = QueryableMethods.AnyWithPredicate.MakeGenericMethod(typeArgument);

            var anyLambdaParameter = Expression.Parameter(typeArgument, "p");
            var anyLambda = Expression.Lambda(
                Expression.Equal(
                    anyLambdaParameter,
                    methodCallExpression.Arguments[1]),
                anyLambdaParameter);

            return Expression.Call(null, anyMethod, new[] { Visit(methodCallExpression.Arguments[0]), anyLambda });
        }

        var @object = default(Expression);
        if (methodCallExpression.Object != null)
        {
            @object = MatchExpressionType(
                Visit(methodCallExpression.Object), methodCallExpression.Object.Type);
        }

        var arguments = new Expression[methodCallExpression.Arguments.Count];
        for (var i = 0; i < arguments.Length; i++)
        {
            arguments[i] = MatchExpressionType(
                Visit(methodCallExpression.Arguments[i]), methodCallExpression.Arguments[i].Type);
        }

        var visited = methodCallExpression.Update(@object!, arguments);

        // In VB.NET, comparison operators between strings (equality, greater-than, less-than) yield
        // calls to a VB-specific CompareString method. Normalize that to string.Compare.
        if (visited.Method.Name == "CompareString"
            && (visited.Method.DeclaringType?.Name == "Operators"
                || visited.Method.DeclaringType?.Name == "EmbeddedOperators")
            && visited.Method.DeclaringType?.Namespace == "Microsoft.VisualBasic.CompilerServices"
            && visited.Object == null
            && visited.Arguments.Count == 3
            && visited.Arguments[2] is ConstantExpression textCompareConstantExpression)
        {
            return textCompareConstantExpression.Value is bool boolValue
                && boolValue
                    ? Expression.Call(
                        StringCompareWithComparisonMethod,
                        visited.Arguments[0],
                        visited.Arguments[1],
                        Expression.Constant(StringComparison.OrdinalIgnoreCase))
                    : Expression.Call(
                        StringCompareWithoutComparisonMethod,
                        visited.Arguments[0],
                        visited.Arguments[1]);
        }

        return visited;
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitNew(NewExpression newExpression)
    {
        if (newExpression.Arguments.Count == 0)
        {
            return newExpression;
        }

        var arguments = new Expression[newExpression.Arguments.Count];
        for (var i = 0; i < arguments.Length; i++)
        {
            arguments[i] = MatchExpressionType(
                Visit(newExpression.Arguments[i]),
                newExpression.Arguments[i].Type);
        }

        return newExpression.Update(arguments);
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitNewArray(NewArrayExpression newArrayExpression)
    {
        var expressions = new Expression[newArrayExpression.Expressions.Count];
        for (var i = 0; i < expressions.Length; i++)
        {
            expressions[i] = MatchExpressionType(
                Visit(newArrayExpression.Expressions[i]),
                newArrayExpression.Expressions[i].Type);
        }

        return newArrayExpression.Update(expressions);
    }

    /// <summary>
    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
    ///     any release. You should only use it directly in your code with extreme caution and knowing that
    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
    /// </summary>
    protected override Expression VisitUnary(UnaryExpression unaryExpression)
    {
        if (unaryExpression.NodeType == ExpressionType.Not
            && unaryExpression.Operand is MethodCallExpression innerMethodCall
            && (Equals(StartsWithMethodInfo, innerMethodCall.Method)
                || Equals(EndsWithMethodInfo, innerMethodCall.Method)))
        {
            if (innerMethodCall.Arguments[0] is ConstantExpression constantArgument
                && constantArgument.Value is string stringValue
                && stringValue == string.Empty)
            {
                // every string starts/ends with empty string.
                return Expression.Constant(false);
            }

            var newObject = Visit(innerMethodCall.Object)!;
            var newArgument = Visit(innerMethodCall.Arguments[0]);

            var result = Expression.AndAlso(
                Expression.NotEqual(newObject, ConstantNullString),
                Expression.AndAlso(
                    Expression.NotEqual(newArgument, ConstantNullString),
                    Expression.Not(innerMethodCall.Update(newObject, new[] { newArgument }))));

            return newArgument is ConstantExpression
                ? result
                : Expression.AndAlso(
                    Expression.NotEqual(
                        newArgument,
                        Expression.Constant(string.Empty)),
                    result);
        }

        return unaryExpression.Update(
            Visit(unaryExpression.Operand));
    }

    private static Expression MatchExpressionType(Expression expression, Type typeToMatch)
        => expression.Type != typeToMatch
            ? Expression.Convert(expression, typeToMatch)
            : expression;

    private static bool TryExtractEqualityOperands(
        Expression expression,
        [NotNullWhen(true)] out Expression? left,
        [NotNullWhen(true)] out Expression? right,
        out bool negated)
    {
        (left, right, negated) = (default, default, default);

        switch (expression)
        {
            case BinaryExpression binaryExpression:
                switch (binaryExpression.NodeType)
                {
                    case ExpressionType.Equal:
                        negated = false;
                        break;
                    case ExpressionType.NotEqual:
                        negated = true;
                        break;
                    default:
                        return false;
                }

                (left, right) = (binaryExpression.Left, binaryExpression.Right);
                return true;

            case MethodCallExpression methodCallExpression
                when methodCallExpression.Method.Name == nameof(object.Equals):
            {
                negated = false;
                if (methodCallExpression.Arguments.Count == 1
                    && methodCallExpression.Object?.Type == methodCallExpression.Arguments[0].Type)
                {
                    (left, right) = (methodCallExpression.Object, methodCallExpression.Arguments[0]);

                    return true;
                }

                if (methodCallExpression.Arguments.Count == 2
                    && methodCallExpression.Arguments[0].Type == methodCallExpression.Arguments[1].Type)
                {
                    (left, right) = (methodCallExpression.Arguments[0], methodCallExpression.Arguments[1]);

                    return true;
                }

                return false;
            }

            case UnaryExpression unaryExpression
                when unaryExpression.IsLogicalNot():
            {
                var result = TryExtractEqualityOperands(unaryExpression.Operand, out left, out right, out negated);
                negated = !negated;
                return result;
            }
        }

        return false;
    }

    private static Expression? TryOptimizeMemberAccessOverConditional(Expression expression)
    {
        // Simplify (a != null ? new { Member = b, ... } : null).Member
        // to a != null ? b : null
        // Later null check removal will simplify it further
        if (expression is MemberExpression visitedMemberExpression
            && visitedMemberExpression.Expression is ConditionalExpression conditionalExpression
            && conditionalExpression.Test is BinaryExpression binaryTest
            && (binaryTest.NodeType == ExpressionType.Equal
                || binaryTest.NodeType == ExpressionType.NotEqual)
            // Exclude HasValue/Value over Nullable<> as they return non-null type and we don't have equivalent for it for null part
            && !(conditionalExpression.Type.IsNullableValueType()
                && (visitedMemberExpression.Member.Name == nameof(Nullable<int>.HasValue)
                    || visitedMemberExpression.Member.Name == nameof(Nullable<int>.Value))))
        {
            var isLeftNullConstant = IsNullConstant(binaryTest.Left);
            var isRightNullConstant = IsNullConstant(binaryTest.Right);

            if (isLeftNullConstant != isRightNullConstant
                && ((binaryTest.NodeType == ExpressionType.Equal
                        && IsNullConstant(conditionalExpression.IfTrue))
                    || (binaryTest.NodeType == ExpressionType.NotEqual
                        && IsNullConstant(conditionalExpression.IfFalse))))
            {
                var nonNullExpression = binaryTest.NodeType == ExpressionType.Equal
                    ? conditionalExpression.IfFalse
                    : conditionalExpression.IfTrue;

                // Use ReplacingExpressionVisitor rather than creating MemberExpression
                // So that member access chain on NewExpression/MemberInitExpression condenses
                nonNullExpression = ReplacingExpressionVisitor.Replace(
                    visitedMemberExpression.Expression, nonNullExpression, visitedMemberExpression);
                nonNullExpression = TryOptimizeMemberAccessOverConditional(nonNullExpression) ?? nonNullExpression;
                if (!nonNullExpression.Type.IsNullableType())
                {
                    nonNullExpression = Expression.Convert(nonNullExpression, nonNullExpression.Type.MakeNullable());
                }

                var nullExpression = Expression.Constant(null, nonNullExpression.Type);

                return Expression.Condition(
                    conditionalExpression.Test,
                    binaryTest.NodeType == ExpressionType.Equal ? nullExpression : nonNullExpression,
                    binaryTest.NodeType == ExpressionType.Equal ? nonNullExpression : nullExpression);
            }
        }

        return null;
    }

    private static bool IsNullConstant(Expression expression)
        => expression is ConstantExpression constantExpression
            && constantExpression.Value == null;
}
